using Microsoft.SqlServer.TransactSql.ScriptDom;

namespace SPI.AgentServer.Services.Security;

public class SqlParser : ISqlParser
{
    private readonly ILogger<SqlParser> _logger;

    public SqlParser(ILogger<SqlParser> logger)
    {
        _logger = logger;
    }

    public SqlParseResult Parse(string sql)
    {
        var result = new SqlParseResult { Success = true };

        try
        {
            var parser = new TSql160Parser(true);
            var fragment = parser.Parse(new StringReader(sql), out var errors);

            if (errors.Count > 0)
            {
                result.Success = false;
                result.Error = string.Join("; ", errors.Select(e => e.Message));
                return result;
            }

            var visitor = new SqlVisitor();
            fragment.Accept(visitor);

            result.StatementType = visitor.StatementType;
            result.Tables = visitor.Tables.Distinct().ToList();
            result.Columns = visitor.Columns.Distinct().ToList();
            result.Operations = visitor.Operations.Distinct().ToList();

            return result;
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error parsing SQL");
            result.Success = false;
            result.Error = $"Parse error: {ex.Message}";
            return result;
        }
    }

    private class SqlVisitor : TSqlFragmentVisitor
    {
        public SqlStatementType StatementType { get; private set; } = SqlStatementType.Other;
        public HashSet<string> Tables { get; } = new();
        public HashSet<string> Columns { get; } = new();
        public HashSet<string> Operations { get; } = new();

        public override void Visit(SelectStatement node)
        {
            StatementType = SqlStatementType.Select;
            Operations.Add("SELECT");
            base.Visit(node);
        }

        public override void Visit(InsertStatement node)
        {
            StatementType = SqlStatementType.Insert;
            Operations.Add("INSERT");
            base.Visit(node);
        }

        public override void Visit(UpdateStatement node)
        {
            StatementType = SqlStatementType.Update;
            Operations.Add("UPDATE");
            base.Visit(node);
        }

        public override void Visit(DeleteStatement node)
        {
            StatementType = SqlStatementType.Delete;
            Operations.Add("DELETE");
            base.Visit(node);
        }

        public override void Visit(CreateTableStatement node)
        {
            StatementType = SqlStatementType.Create;
            Operations.Add("CREATE");
            base.Visit(node);
        }

        public override void Visit(AlterTableStatement node)
        {
            StatementType = SqlStatementType.Alter;
            Operations.Add("ALTER");
            base.Visit(node);
        }

        public override void Visit(DropTableStatement node)
        {
            StatementType = SqlStatementType.Drop;
            Operations.Add("DROP");
            base.Visit(node);
        }

        public override void Visit(TruncateTableStatement node)
        {
            StatementType = SqlStatementType.Truncate;
            Operations.Add("TRUNCATE");
            base.Visit(node);
        }

        public override void Visit(NamedTableReference node)
        {
            var tableName = node.SchemaObject.BaseIdentifier.Value;
            if (node.SchemaObject.SchemaIdentifier != null)
            {
                tableName = $"{node.SchemaObject.SchemaIdentifier.Value}.{tableName}";
            }
            Tables.Add(tableName);
            base.Visit(node);
        }

        public override void Visit(ColumnReferenceExpression node)
        {
            if (node.MultiPartIdentifier != null)
            {
                var columnName = string.Join(".", node.MultiPartIdentifier.Identifiers.Select(i => i.Value));
                Columns.Add(columnName);
            }
            base.Visit(node);
        }
    }
}
